## visualizing chaotic tiemseries by plotly and hypertools

import numpy as np
import pandas as pd
from scipy.integrate import solve_ivp
import matplotlib.pyplot as plt

import plotly
import plotly.express as px
pd.options.plotting.backend = "plotly"
plotly.offline.init_notebook_mode() # plotly 擅長交互

# import hypertools as hyp
# 需要在 Jupyter Notebook 环境下，才能实现动态效果
# 由另一個單元打開交互圖時，動圖就會停止
# %matplotlib nbAgg
# %matplotlib ipympl # 只適合於交互，

def lorenz_deriv(t, x_y_z, a, b, c):
    """定义所要求解的 ODE，计算 Lorenz 系统的时间导数"""
    x, y, z = x_y_z
    return [a*(y - x), c*x - x*z - y, x*y - b*z]

def chen_deriv(t, x_y_z, a, b, c):
    """定义所要求解的 ODE，计算 Chen 系统的时间导数"""
    x, y, z = x_y_z
    return [a*(y - x), (c - a)*x - x*z + c*y, x*y - b*z]

def my_ani_data(data, rl=1):
    """
    : data, array of 3d line corrd.
    : rl, length of each repeat.
    Suggested ratio is len(data)/rl <= 100
    """
    x_t_df = pd.DataFrame(data, columns=['x', 'y', 'z'])
    rn = len(x_t_df)//rl+1 # number of repeats
    x_t_df['ani_frame'] = np.repeat(np.arange(rn), rl)[:len(x_t_df)]
    x_t_df_cum = pd.DataFrame(columns=['x', 'y', 'z', 'ani_frame'])
    x_t_df_new = pd.DataFrame(columns=['x', 'y', 'z', 'ani_frame'])
    for i in range(rn):
        x_t_df_i = x_t_df[x_t_df['ani_frame']==i] 
        x_t_df_cum = pd.concat([x_t_df_cum, x_t_df_i])
        x_t_df_cum['ani_frame'] = i
        x_t_df_new = pd.concat([x_t_df_new, x_t_df_cum])
    return x_t_df_new

## 示例 0：rk4 和 lsoda 求解器的對比

max_pt = 20000
tau = 0.05 # 延迟时间
t_max = max_pt * tau
t = np.linspace(0, t_max, num=max_pt)

# help(solve_ivp)

# method='RK45'
lorenz_sol = solve_ivp(fun=lorenz_deriv, t_span=[0, t_max], y0=[1, 0, 1], t_eval=t, args=(10,8/3,28), dense_output=False)
lorenz_sol_df = pd.DataFrame(lorenz_sol.y.T, index=lorenz_sol.t)
chen_sol = solve_ivp(fun=chen_deriv, t_span=[0, t_max], y0=[-1, -1.1, 0], t_eval=t, args=(35,3,28), dense_output=False)
chen_sol_df = pd.DataFrame(chen_sol.y.T, index=chen_sol.t)
# lorenz_sol_df.to_csv('lorenz_rk45.csv')
# chen_sol_df.to_csv('chen_rk45.csv')
# lorenz_rk4 = pd.read_csv('lorenz_rk45.csv', index_col=0)
# chen_rk4 = pd.read_csv('chen_rk45.csv', index_col=0)
# hyp.plot(hyp.tools.df2mat(lorenz_rk4), animate=True)
# hyp.plot(hyp.tools.df2mat(chen_rk4), animate=True)

# method='LSODA'
lorenz_sol = solve_ivp(fun=lorenz_deriv, t_span=[0, t_max], y0=[1, 0, 1], method='LSODA', t_eval=t, args=(10,8/3,28))
lorenz_sol_df = pd.DataFrame(lorenz_sol.y.T, index=lorenz_sol.t)
chen_sol = solve_ivp(fun=chen_deriv, t_span=[0, t_max], y0=[-1, -1.1, 0], method='LSODA', t_eval=t, args=(35,3,28))
chen_sol_df = pd.DataFrame(chen_sol.y.T, index=chen_sol.t)
# lorenz_sol_df.to_csv('lorenz_lsoda.csv')
# chen_sol_df.to_csv('chen_lsoda.csv')
# lorenz_lsoda = pd.read_csv('lorenz_lsoda.csv', index_col=0)
# chen_lsoda = pd.read_csv('chen_lsoda.csv', index_col=0)
# hyp.plot(hyp.tools.df2mat(lorenz_lsoda), animate=True)
# hyp.plot(hyp.tools.df2mat(chen_lsoda), animate=True)

## 示例 1：Lorenz 系統的可視化
## 相空間中粒子的運動軌跡（吸引子）呈蝴蝶形；
## 運動十分不穩定，例如，軌道在右邊轉一圈，然後去到左邊轉幾圈，再回到右邊轉等等；

max_pt = 20000 + 20000 + 2000 + 2000
tau = 0.05 # 延迟时间 0.01 * np.array([1,3,5,7,10])
t_max = max_pt * tau
t = np.linspace(0, t_max, num=max_pt)

# method='RK45' by default
lorenz_sol = solve_ivp(fun=lorenz_deriv, t_span=[0, t_max], y0=[1, 0, 1], t_eval=t, args=(10,8/3,28))
lorenz_sol_df = pd.DataFrame(lorenz_sol.y.T, index=lorenz_sol.t)

# lorenz_sol_df.iloc[20000:].to_csv('lorenz_XYC005.csv') # 捨棄瞬態
# lorenz_sol_df = pd.read_csv('lorenz_XYC005.csv', index_col=0)

# 時序圖
lorenz_sol_df.iloc[20000:].plot()

# 相圖
hyp.plot(hyp.tools.df2mat(lorenz_sol_df.iloc[20000:]), 
         animate=True,
         duration=60,
         chemtrails=True,
        )
# help(hyp.plot)

## 示例 2：Chen 系統的可視化
## 混沌吸引子軌道對初始條件有著極端的敏感性，
## 即取兩個鄰近的初始點，它們的軌道是以指數速度分離而且轉的圈數也不同。

max_pt = 500
tau = 0.05 # 延迟时间
t_max = max_pt * tau
t = np.linspace(0, t_max, num=max_pt)

# method='RK45'
chen_sol = solve_ivp(fun=chen_deriv, t_span=[0, t_max], y0=[-1, -1.1, 12], t_eval=t, args=(35,3,28))
chen_sol_df0 = pd.DataFrame(chen_sol.y.T, index=chen_sol.t)
chen_sol = solve_ivp(fun=chen_deriv, t_span=[0, t_max], y0=[-1, -1.1, 12+1e-4], t_eval=t, args=(35,3,28))
chen_sol_df1 = pd.DataFrame(chen_sol.y.T, index=chen_sol.t)

# 時序圖
chen_sol_df0.columns = ['x0','y0','z0']
chen_sol_df1.columns = ['x1','y1','z1']
_ = pd.concat([chen_sol_df0,chen_sol_df1],axis=1).plot()
_.show()

# 相圖
ani_data0 = my_ani_data(chen_sol_df0.values, rl=10)
ani_data1 = my_ani_data(chen_sol_df1.values, rl=10)
ani_data0['label'] = 'chen0'
ani_data1['label'] = 'chen1'
fig = px.line_3d(pd.concat([ani_data0,ani_data1]), x='x', y='y', z='z', animation_frame='ani_frame', color='label')
fig.update_layout(autosize=False, width=800, height=800)
fig.show()


## refer
# https://hypertools.readthedocs.io/en/latest/
# https://plotly.com/python/3d-line-plots/
# https://plotly.com/python/animations/

